package org.chalmers.rl;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Queue;

import com.google.common.collect.EvictingQueue;

public class StateAction {
	public static final int MAX_MEMORY = 100;
	
	private double val = 0;
	public double avg_r = 0;
	
	public Integer myAct;
	public Integer myState;
	
	
	private double lastSampledReward;
	private int lastSampledState;
	
	
	//private Map<Integer, Integer> stateTransitions = new HashMap<Integer, Integer>();
	//private Map<Integer, Integer> lastTime = new HashMap<Integer, Integer>();
	
	private EvictingQueue<Double> rewards = EvictingQueue.create(MAX_MEMORY);
	private EvictingQueue<Integer> states = EvictingQueue.create(MAX_MEMORY);
	
	private Map<Integer, Double> state_total_rewards = new HashMap<Integer, Double>();
	private Map<Integer, Integer> state_count = new HashMap<Integer, Integer>();
	
	
	
	public StateAction(Integer actIndex, Integer stateIndex, double initVal) {
		this.myAct = actIndex;
		this.myState = stateIndex;
		this.val = initVal;
	}
	
	public void observe(Integer toState, double reward) {
		//step_n++;
		//alpha = 1.0/step_n; //simple averaging for rewards;
		//avg_r += (reward - avg_r)*alpha; 
		
		//lastTime.put(toState, step_n);
//		if (states.size() == MAX_MEMORY) {
//			Integer evict = states.peek();
//			Double rew = rewards.peek();
//			state_total_rewards.put(evict, state_total_rewards.get(evict) - rew);
//			state_count.put(evict, state_count.get(evict) - 1);
//		}
		
		rewards.add(reward);
		states.add(toState);
		
		if (!state_count.containsKey(toState)) {
			state_count.put(toState, 1);
			state_total_rewards.put(toState, reward);
			return;
		}
		
		state_count.put(toState, state_count.get(toState) + 1);
		int count = state_count.get(toState);
		double rew = state_total_rewards.get(toState);
		
		rew = rew + reward;
		state_total_rewards.put(toState,rew);

	}
	
//	public Set<Integer> accessibleStates() {
//		return stateTransitions.keySet();
//	}
//	
//	public double trans_p(Integer s) {
//		if (stateTransitions.containsKey(s)) {
//			return (1.0*stateTransitions.get(s))/step_n;
//		}
//		return 0;
//	}
	
//	public void end_episode() {
//		step_n = 0;
//	}
//	public double avgRewardForTransition(Integer nextState) {
//		double totalR = 0;
//		int count = 0;
//		Iterator<Integer> sIterator = states.iterator();
//		Iterator<Double> rIterator = rewards.iterator();
//		while (sIterator.hasNext()) { //go along states and rewards simultaneously
//			Integer s = sIterator.next();
//			Double r = rIterator.next();
//			if (s.equals(nextState)) {
//				count++;
//				totalR += r;
//			}
//		} 
//		if (count == 0) {
//			return 0;
//		} else {
//			return totalR/count;
//		}
//	} 
	
	public double avgRewardForTransition(Integer nextState) {
		return !state_total_rewards.containsKey(nextState) ? 0 : 
			state_total_rewards.get(nextState)/state_count.get(nextState);
	} 
	
	
	public Queue<Integer> transitions() {
		return states;
	}
	
	public Queue<Double> rewards() {
		return rewards;
	}
	
	public double lastSampledReward() {
		return lastSampledReward;
	}
	
	public int lastSampledState() {
		return lastSampledState;
	}
	
	// a bit ugly way to sample reward and state
	public void sample() {
		int randIndex = Utils.rand.nextInt(states.size());
		lastSampledReward = (Double) rewards.toArray()[randIndex];
		lastSampledState = (Integer) states.toArray()[randIndex];
	}
	
	public String toString() {
		StringBuffer sb = new StringBuffer();
		sb.append("My state: " + myState)
		  .append(" action: " + myAct)
		  .append(" val " + val)
		  .append(" Possibe transitions: ");
		
		Iterator<Integer> sIterator = states.iterator();
		Iterator<Double> rIterator = rewards.iterator();
		while (sIterator.hasNext()) { //go along states and rewards simultaneously
			Integer s = sIterator.next();
			Double r = rIterator.next();
			sb.append("s: " + s + " r: " + r + ", ");
		}
		
		return sb.toString();
		
	}

	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		result = prime * result + ((myAct == null) ? 0 : myAct.hashCode());
		result = prime * result + ((myState == null) ? 0 : myState.hashCode());
		return result;
	}

	@Override
	public boolean equals(Object obj) {
		if (this == obj)
			return true;
		if (obj == null)
			return false;
		if (getClass() != obj.getClass())
			return false;
		StateAction other = (StateAction) obj;
		if (myAct == null) {
			if (other.myAct != null)
				return false;
		} else if (!myAct.equals(other.myAct))
			return false;
		if (myState == null) {
			if (other.myState != null)
				return false;
		} else if (!myState.equals(other.myState))
			return false;
		return true;
	}

	public double val() {
		return val;
	}
	
	public void setVal(double v) {
		this.val = v;
	}
}
